#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot shallow water Rossby-Haurwitz wave test results')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

f = xr.open_dataset('swm_rh.720x360.h0.nc')

z_levels = np.linspace(8000, 10500, 13)

fig = plt.figure(figsize=(10, 5))
proj = ccrs.PlateCarree()
ax = plt.axes(projection=proj)
plot_contour_map(ax, f.gz[args.time_step,0,:,:] / 9.80616, cmap='jet', left_string=f'Geopotential height (m) at {((f.time[args.time_step] - f.time[0]).values).astype("timedelta64[D]")}', levels=z_levels)

if args.output:
	plt.savefig(args.output, format='pdf')
else:
	plt.show()
